import paddle
import paddle.fluid as fluid

__all__ = [
    "ResNet"
]

from ppcls.modeling.architectures.cifar.layers import conv2d, pool2d, fc


def shortcut(x, out_channels, stride, name):
    in_channels = x.shape[1]
    if in_channels != out_channels or stride != 1:
        if stride != 1:
            x = pool2d(x, 2, stride=2, type='avg', name = name + '.pool')
        x = conv2d(x, out_channels, 1, 1, name=name + '.conv')
    return x


def basic_block(x, out_channels, stride, name):
    identity = x
    x = conv2d(x, out_channels, kernel_size=3, stride=stride,
               bn=True, act='relu',  name=name + ".branch2a")
    x = conv2d(x, out_channels, kernel_size=3,
               bn=True, act='relu',  name=name + ".branch2b")
    identity = shortcut(identity, out_channels, stride, name=name + ".branch1")
    return fluid.layers.elementwise_add(x=identity, y=x, act='relu')


def bottleneck(x, out_channels, stride, name):
    identity = x
    channels = out_channels // 4
    x = conv2d(x, channels, kernel_size=1,
                   bn=True, act='relu', name=name + ".branch2a")
    x = conv2d(x, channels, kernel_size=3,
                   bn=True, act='relu', name=name + ".branch2b")
    x = conv2d(x, out_channels, kernel_size=1,
                   bn=True, act=None, name=name + ".branch2c")
    identity = shortcut(identity, out_channels, stride, name=name + ".branch1")
    return fluid.layers.elementwise_add(x=identity, y=x, act='relu')


class ResNet:

    def __init__(self, depth, block):
        self.depth = depth
        self.block = block

    def net(self, input, class_dim=10):
        depth = self.depth

        if self.block == 'basic':
            block = basic_block
            layers = [(depth - 2) // 6] * 3
        else:
            block = bottleneck
            layers = [(depth - 2) // 9] * 3

        stages = [16, 32, 64]

        x = conv2d(input, 16, kernel_size=3, bn=True, act='relu')

        x = self._make_layer(x, block, stages[0], layers[0], stride=1, name='stage1')
        x = self._make_layer(x, block, stages[1], layers[1], stride=2, name='stage2')
        x = self._make_layer(x, block, stages[2], layers[2], stride=2, name='stage3')

        x = fluid.layers.pool2d(input=x, pool_type='avg', global_pooling=True)

        x = fc(x, class_dim)
        return x

    def _make_layer(self, x, block, out_channels, num_units, stride, name):
        x = block(x, out_channels, stride=stride, name=name + ".unit1")
        for i in range(1, num_units):
            x = block(x, out_channels, stride=1, name=name + f".unit{i+1}")
        return x